import torch
from torch import nn
from dataset import myDataSet
from torch.utils.data import DataLoader
from u2net import U2NETP
import os
from torchvision.utils import save_image
from torchvision import transforms
from PIL import Image
from tools import imgScale
import numpy as np

if __name__ == '__main__':

    # test_img_path = r'D:\project\opintImg\xiangzi'
    # test_img_path = r'./data/pictiure/'
    test_img_path = r'./data/pics_at_night/'

    test_file_name = 'stream0_1620743650.59416.jpg'

    netSavePath = './params/145_0.058549691722030714.pth'

    test_path = r'./data/test_imgs/'

    img = Image.open(os.path.join(test_img_path, test_file_name))

    # 数据预处理
    # img = imgScale(img, 320)

    img_data = transforms.ToTensor()(img)

    # img_data = transforms.Compose([
    #             transforms.ToTensor(),
    #             # transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    #             ])(img)

    # 增加一维
    img_data = img_data.unsqueeze(0)

    # 设定运行网络
    net = U2NETP()
    # 加载参数
    if os.path.exists(netSavePath):
        net.load_state_dict(torch.load(netSavePath, map_location="cuda:0"))

        print('u2net神经网络参数加载成功！')
    else:
        print('u2net神经网络重新创建连接！')

    # 将程序放在网络中
    d0, d1, d2, d3, d4, d5, d6 = net(img_data)

    # # 制作mask
    mask = d0[0].detach().numpy()
    mask = np.where(mask > 0.5, 1, 0).transpose(1, 2, 0).astype(np.uint8)

    # 结果输出
    out_Img = torch.cat([d0[0] * 255, d0[0] * 255, d0[0] * 255], 0)
    out_Img = out_Img.detach().numpy()

    out_Img = out_Img.transpose(1, 2, 0).astype(np.uint8)

    # img2 = Image.fromarray(img * mask)
    img2 = Image.fromarray(out_Img)
    # img2.show()
    # exit()
    img2.save(os.path.join(test_path, test_file_name))
